#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -*- python -*-
"""
 gen-primetab.py

 by Jacob Smith <http://code.google.com/p/lua-mpi/>
 Copyright (C) 2009 Jacob Smith, All Rights Reserved

 $Id: gen-primetab.py 16 2009-08-06 04:15:21Z kloptops $
"""

import math;
from util import *
from optparse import OptionParser

def gen_primes(n):
    """
    gen_primes(n) --> prime_list

    Return list of primes from 2 up to but not including n.  Uses Sieve of Erasth.
    """
    if n < 2:
        return []

    nums = range(2, int(n) + 1)
    num_hits = [0 for i in nums];
    p = []
    x = 0;

    while x < len(nums):
        if num_hits[x] > 0:
            x += 1;
            continue;

        new_prime = nums[x]
        p.append(new_prime)
        x += 1;

        for i in xrange(x, len(nums)):
            if num_hits[i] == 0 and nums[i] % new_prime == 0:
                num_hits[i] += 1;

    print "gen_primes(%d) = %d" % (n, len(p));
    return p;


def compress(nums, c_map):
    c_map_c = c_map[0]
    c_map_x = c_map[1]
    c_map_l = len(c_map[0])

    result = "";
    last = 0;
    for i in nums:
        diff = i - last - 1;
        last = i;
        while diff >= c_map_l:
            result += c_map_x;
            diff -= c_map_l;

        result += c_map_c[diff];

    return result

compress_map = ["0123456789ABCDEF", "/"];


####
# Compressed Output.
def gen_c_output(c_primes, n_primes, compress_map):
    num_per_line = 80 - 1 - 4;
    full = len(n_primes);
    half = (full + 1) / 2;

    c_output  = 'local s_mpi_primetab = {}; -- Decompresses at runtime...\n';
    c_output += 'do local t, m, mc, ml = s_mpi_primetab, "' + compress_map[0] + '", "' + compress_map[1] + '", ' + str(len(compress_map[0])) + ';\n';
    c_output += 'local e, i, l, p = 0, 1, 0, "" ..\n';
    for i in xrange(0, len(c_primes), num_per_line):
        c_output += '"' + c_primes[(i):(i + num_per_line)] + '"' + (i + num_per_line < len(c_primes) and " .." or ";") + '\n';
    c_output += 'while i <= string.len(p) do e = string.sub(p, i, i); i = i + 1;\n';
    c_output += 'while e == mc do l = l + ml; e = string.sub(p, i, i); i = i + 1; end;\n';
    c_output += 'l = l + string.find(m, e, 0, true); t[#t + 1] = l; end;\n';
    c_output += 'assert(t[' + str(half) + '] == ' + str(n_primes[half - 1]) + ' and t[' + str(full) + '] == ' + str(n_primes[full - 1]) + ',\n';
    c_output += '("mpi_primetab: %d ~= ' + str(n_primes[half - 1]) + ' or %d ~= ' + str(n_primes[full - 1]) + '"):format(t[' + str(half) + '], t[' + str(full) + ']))\n';
    c_output += 'end\n';

    return c_output;

####
# Non-Compressed Output.
def gen_n_output(n_primes):
    pack_per_line = 10;

    n_output  = "local s_mpi_primetab = {\n";
    n_output += "\n".join(map(lambda ix: ", ".join(map(lambda x: "0x%04x" % (x), n_primes[ix:(ix + pack_per_line)])) + ", ", range(0, len(n_primes), pack_per_line)));
    n_output += "}";

    return n_output;

############################

def main():
    parser = OptionParser(usage = "usage: %prog [options] [32 | 64 | 128 | 256 | 512 | 1024 | 8bit | 16bit]")
    parser.add_option("-n", "--no-compress",
                      action="store_false", dest="do_compress", default=True,
                      help="don't compress prime list.");

    parser.add_option("-p", "--print",
                      action="store_true", dest="do_print", default=False,
                      help="print out generated code, then exit.");

    (options, args) = parser.parse_args()

    if len(args) == 0:
        args = ['16bit']

    max_primes = {
        '32':    0x0035,
        '64':    0x0137,
        '128':   0x02CF,
        '256':   0x0653,
        '512':   0x0E59,
        '1024':  0x1FE7,
        '8bit':    0xFF,
        '16bit': 0xFFFF,
        };

    if not max_primes.has_key(args[0]):
        parser.error("invalid option '" + args[0] + "' choices are 32, 64, 128, 256, 512, 1024, 8bit, 16bit.");

    list_size = max_primes[args[0]];

    n_primes = gen_primes(list_size);
    c_primes = compress(n_primes, compress_map);


    c_output = gen_c_output(c_primes, n_primes, compress_map);
    n_output = gen_n_output(n_primes);

    output = '';

    if options.do_compress:
        print len(n_output) - len(c_output)
        if (len(n_output) - len(c_output)) > 512:
            output = c_output;
        else:
            output = n_output;
    else:
        output = n_output;

    ####
    # Output...

    mpi_file = abspath('mpi.lua')
    if not isfile(mpi_file) or options.do_print:
        print len(c_output), len(n_output);
        print "-" * 80;
        print n_output;
        print "-" * 80;
        print c_output;
    else:
        zone = Zones()

        data = loadraw(mpi_file);
        zone.unpack(data);
        zone.replace('PRIME TAB', output);
        packed = zone.pack();
        saveraw(mpi_file, packed);

if __name__ == '__main__':
    main();

